package database

import (
	"database/sql"
	"fmt"
	"log"
	"stock-analyzer/internal/model"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

// DB 数据库连接
type DB struct {
	conn *sql.DB
}

// NewDB 创建新的数据库连接
func NewDB(dbPath string) (*DB, error) {
	conn, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}

	// 创建表结构
	if err := createTable(conn); err != nil {
		return nil, err
	}

	return &DB{conn: conn}, nil
}

// createTable 创建股票数据表
func createTable(conn *sql.DB) error {
	createTableSQL := `
	CREATE TABLE IF NOT EXISTS stock_data (
		stock_code TEXT NOT NULL,
		date TEXT NOT NULL,
		open REAL,
		high REAL,
		low REAL,
		close REAL,
		volume INTEGER,
		amount REAL,
		PRIMARY KEY (stock_code, date)
	)
	`

	_, err := conn.Exec(createTableSQL)
	if err != nil {
		return fmt.Errorf("failed to create table: %w", err)
	}

	return nil
}

// Close 关闭数据库连接
func (db *DB) Close() error {
	if db.conn != nil {
		return db.conn.Close()
	}
	return nil
}

// GetStockDataByTimeRange 获取指定时间范围内的股票数据
func (db *DB) GetStockDataByTimeRange(stockCode string, startDate, endDate time.Time) ([]model.StockData, error) {
	query := `
		SELECT stock_code, date, open, high, low, close, volume, amount
		FROM stock_data
		WHERE stock_code = ? AND date BETWEEN ? AND ?
		ORDER BY date ASC
	`

	rows, err := db.conn.Query(query, stockCode, startDate.Format("20060102"), endDate.Format("20060102"))
	if err != nil {
		return nil, fmt.Errorf("failed to query stock data: %w", err)
	}
	defer rows.Close()

	var results []model.StockData
	for rows.Next() {
		var data model.StockData
		var dateStr string

		err := rows.Scan(
			&data.StockCode,
			&dateStr,
			&data.Open,
			&data.High,
			&data.Low,
			&data.Close,
			&data.Volume,
			&data.Amount,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan row: %w", err)
		}

		// 解析日期字符串
		data.Date, err = time.Parse("20060102", dateStr)
		if err != nil {
			log.Printf("Warning: failed to parse date %s: %v", dateStr, err)
			continue
		}

		results = append(results, data)
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("error iterating rows: %w", err)
	}

	return results, nil
}

// GetAllStockCodes 获取所有股票代码
func (db *DB) GetAllStockCodes() ([]string, error) {
	query := `SELECT DISTINCT stock_code FROM stock_data ORDER BY stock_code`
	rows, err := db.conn.Query(query)
	if err != nil {
		return nil, fmt.Errorf("failed to query stock codes: %w", err)
	}
	defer rows.Close()

	var stockCodes []string
	for rows.Next() {
		var code string
		if err := rows.Scan(&code); err != nil {
			return nil, fmt.Errorf("failed to scan stock code: %w", err)
		}
		stockCodes = append(stockCodes, code)
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("error iterating stock codes: %w", err)
	}

	return stockCodes, nil
}

// InsertStockData 插入股票数据
func (db *DB) InsertStockData(data model.StockData) error {
	insertSQL := `
		INSERT OR REPLACE INTO stock_data 
		(stock_code, date, open, high, low, close, volume, amount)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?)
	`

	_, err := db.conn.Exec(
		insertSQL,
		data.StockCode,
		data.Date.Format("2006-01-02"),
		data.Open,
		data.High,
		data.Low,
		data.Close,
		data.Volume,
		data.Amount,
	)

	if err != nil {
		return fmt.Errorf("failed to insert stock data: %w", err)
	}

	return nil
}
